Fix KTO evaluate crash on streaming datasets with unpaired-preference data by albertvillanova · Pull Request #6325 · huggingface/trl

albertvillanova · 2026-07-08T09:53:55Z

This PR fixes streaming evaluation for KTOTrainer with unpaired-preference data, where a batch field returned as a plain list broke Accelerate's dispatch mechanism for IterableDataset.

See: https://github.com/huggingface/trl/actions/runs/28934157283/job/85840257000?pr=6326

FAILED tests/test_kto_trainer.py::TestKTOTrainer::test_evaluate_with_eval_dataset_dict[True] - TypeError: Can only concatenate tensors but got <class 'bool'>

Traceback:

  >       metrics = trainer.evaluate(eval_dataset=eval_dataset)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  
  tests/test_kto_trainer.py:369: 
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  trl/trainer/kto_trainer.py:1648: in evaluate
      return super().evaluate(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4451: in evaluate
      dataset_metrics = self.evaluate(
  trl/trainer/kto_trainer.py:1648: in evaluate
      return super().evaluate(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4469: in evaluate
      output = eval_loop(
  .venv/lib/python3.12/site-packages/transformers/trainer.py:4655: in evaluation_loop
      for step, inputs in enumerate(dataloader):
                          ^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/data_loader.py:856: in __iter__
      next_batch, next_batch_info = self._fetch_batches(main_iterator)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/data_loader.py:812: in _fetch_batches
      batch = concatenate(batches, dim=0)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:617: in concatenate
      return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:615: in concatenate
      return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:81: in honor_type
      return type(obj)(generator)
             ^^^^^^^^^^^^^^^^^^^^
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:615: in <genexpr>
      return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
  
  data = [False], dim = 0
  
      def concatenate(data, dim=0):
          """
          Recursively concatenate the tensors in a nested list/tuple/dictionary of lists of tensors with the same shape.
      
          Args:
              data (nested list/tuple/dictionary of lists of tensors `torch.Tensor`):
                  The data to concatenate.
              dim (`int`, *optional*, defaults to 0):
                  The dimension on which to concatenate.
      
          Returns:
              The same data structure as `data` with all the tensors concatenated.
          """
          if isinstance(data[0], (tuple, list)):
              return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
          elif isinstance(data[0], Mapping):
              return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
          elif not isinstance(data[0], torch.Tensor):
  >           raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
  E           TypeError: Can only concatenate tensors but got <class 'bool'>
  
  .venv/lib/python3.12/site-packages/accelerate/utils/operations.py:619: TypeError

Motivation

DataCollatorForUnpairedPreference and DataCollatorForVisionUnpairedPreference returned the "label" batch field as a plain Python list of booleans instead of a tensor, while every other field in the batch was a tensor. This works fine with a map-style Dataset, but Accelerate's DataLoaderDispatcher, used for streaming IterableDatasets, needs every batch field to be a tensor so it can concatenate and redistribute sub-batches across processes. Passing a streaming eval dataset to KTOTrainer.evaluate() therefore raised TypeError: Can only concatenate tensors but got <class 'bool'>.

Note this is not an Accelerate-version-specific issue. Accelerate ≥1.13.0 added a short-circuit in concatenate() that happens to skip the crash when num_processes == 1 (see huggingface/accelerate#4111, which the Accelerate maintainers confirm is a partial, single-process-only fix). Real multi-process/distributed streaming evaluation still crashes on every Accelerate version, old and new. So this fix is required unconditionally, regardless of installed Accelerate version.

Solution

"label" is now built with torch.tensor(...) in both collators, consistent with every other returned field. Docstrings and doctest examples were updated to reflect the new tensor output.

Changes

Return "label" as a tensor in DataCollatorForUnpairedPreference.torch_call and DataCollatorForVisionUnpairedPreference.torch_call
Update docstrings and doctest examples for both collators to show "label" as a tensor
Update the existing test_padding_and_masks assertion for the new tensor type
Add a regression test that runs a real streaming evaluation through KTOTrainer.evaluate() with an IterableDataset
Add a code comment on both collator changes referencing concatenate() still raises TypeError on non-tensor batch fields when num_processes > 1 accelerate#4111, explaining why this must stay unconditional

Note

Low Risk
Narrow collator output change; loss code already accepts labels via torch.tensor/torch.as_tensor, with regression coverage for streaming eval.

Overview
Fixes KTOTrainer.evaluate() crashing on streaming IterableDataset eval when Accelerate’s dataloader dispatcher tries to concatenate batch fields and hits a plain Python list of booleans for label.

DataCollatorForUnpairedPreference and DataCollatorForVisionUnpairedPreference now build label with torch.tensor(...), matching other batch fields and satisfying Accelerate’s tensor-only concatenation (see accelerate#4111). Docstrings and doctest examples were updated accordingly.

Tests assert tensor labels in test_padding_and_masks and add test_evaluate_with_iterable_dataset for end-to-end streaming evaluation with apo_zero_unpaired (no KL term on streaming data).

^{Reviewed by Cursor Bugbot for commit b268a72. Bugbot is set up for automated code reviews on this repo. Configure here.}

bot-ci-comment · 2026-07-08T09:55:58Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

albertvillanova · 2026-07-09T05:29:39Z

Why this isn't gated on an Accelerate version

While debugging the tests_min_versions CI failure for this fix, I initially assumed it was Accelerate-version-specific, since the regular tests job (latest Accelerate) passes while tests_min_versions (accelerate==1.4.0) fails. Wanted to record the investigation here in case anyone else looks at this and wonders the same thing.

Turns out that's misleading: Accelerate 1.13.0 added a short-circuit in concatenate() that returns a single gathered batch as-is instead of type-checking it.

Allow non-Tensor values in a batch with dispatch_batches=True accelerate#3850

DataLoaderDispatcher._fetch_batches gathers exactly num_processes batches before calling concatenate, so this short-circuit only fires when num_processes == 1, i.e. a plain, non-distributed test run, which is all our CI exercises. It masks the crash for single-process runs on Accelerate ≥1.13.0, but does nothing for num_processes > 1.

Confirmed with the Accelerate maintainers themselves in the PR review thread (huggingface/accelerate#3850 (comment)): they explicitly acknowledged multi-process support was out of scope for that fix

"it'll always fail in multi-process settings... multi-process support is not possible".

I opened a GH issue to track this as a documented, discoverable limitation:

concatenate() still raises TypeError on non-tensor batch fields when num_processes > 1 accelerate#4111

So: real distributed streaming evaluation with KTO's unpaired-preference collator would still crash today, on any Accelerate version, without this fix. The fix here ("label" as a tensor, matching every other collated field) is the correct, version-independent solution, not something to conditionally apply only for old Accelerate.

…bel-tensor

qgallouedec

nice. Plus, it's more consistent to have collators that only return tensors

albertvillanova added 2 commits July 8, 2026 11:50

Return unpaired-preference label as a tensor in KTO collators

fc49e83

Test unpaired-preference label tensor and streaming eval in KTO

efe448b

albertvillanova marked this pull request as draft July 8, 2026 13:17

albertvillanova mentioned this pull request Jul 8, 2026

Support DatasetDict and IterableDatasetDict as eval_dataset in evaluate #6326

Merged

Note the multi-process tensor constraint on KTO label fields

b770291

albertvillanova marked this pull request as ready for review July 9, 2026 05:30

albertvillanova added 3 commits July 9, 2026 07:31

Merge remote-tracking branch 'upstream/main' into fix-kto-collator-la…

7929cf2

…bel-tensor

Merge remote-tracking branch 'upstream/main' into fix-kto-collator-la…

71e6c5f

…bel-tensor

Merge branch 'main' into fix-kto-collator-label-tensor

b268a72

qgallouedec approved these changes Jul 10, 2026

View reviewed changes

albertvillanova merged commit 98ba6a2 into main Jul 11, 2026
13 checks passed

albertvillanova deleted the fix-kto-collator-label-tensor branch July 11, 2026 06:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix KTO evaluate crash on streaming datasets with unpaired-preference data#6325

Fix KTO evaluate crash on streaming datasets with unpaired-preference data#6325
albertvillanova merged 6 commits into
mainfrom
fix-kto-collator-label-tensor

albertvillanova commented Jul 8, 2026 •

edited by cursor Bot

Loading

Uh oh!

bot-ci-comment Bot commented Jul 8, 2026

Uh oh!

albertvillanova commented Jul 9, 2026 •

edited

Loading

Uh oh!

qgallouedec left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

albertvillanova commented Jul 8, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Solution

Changes

Uh oh!

bot-ci-comment Bot commented Jul 8, 2026

Uh oh!

albertvillanova commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why this isn't gated on an Accelerate version

Uh oh!

qgallouedec left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

albertvillanova commented Jul 8, 2026 •

edited by cursor Bot

Loading

albertvillanova commented Jul 9, 2026 •

edited

Loading